use std::fs;
use std::path::PathBuf;

fn main() {
    let out_dir = PathBuf::from("src/onnx");
    let proto_dir = PathBuf::from("proto");

    let protos = [
        proto_dir.join("onnx.proto"),
        // proto_dir.join("onnx-ml.proto"),
    ];

    let mut config = prost_build::Config::new();
    config.out_dir(&out_dir);

    config
        .compile_protos(&protos, &[proto_dir.clone()])
        .expect("expection: failed to compile onnx protos");

    let mut merged = String::new();
    for entry in fs::read_dir(&out_dir).unwrap() {
        let path = entry.unwrap().path();
        if path.extension().and_then(|s| s.to_str()) == Some("rs")
            && path.file_name().and_then(|s| s.to_str()) != Some("mod.rs")
        {
            merged.push_str(&fs::read_to_string(&path).unwrap());
            merged.push('\n');
        }
    }
    fs::write(out_dir.join("onnx.rs"), merged).expect("expection: write merged onnx.rs");
}
